#include "AES.h"
#include "AESCiphereVersion.h"
#include <algorithm>
#include <fstream>
#include <future>
#include <windows.h>

#undef max
#undef min

#pragma comment(lib, "bcrypt.lib")

#ifdef min
    #undef min
#endif
#ifdef max
    #undef max
#endif

using CipherVersion = AES::CipherVersion;

namespace
{
const unsigned int kHeaderSize = 32U;
const unsigned int kMagicSize = 8U;
const char kMagicV3[kMagicSize] = { 'S', 'V', 'I', 'S', 'I', 'O', 'N', '3' };
const char kMagicV2[kMagicSize] = { 'S', 'V', 'I', 'S', 'I', 'O', 'N', '2' };
const char kMagicV1[kMagicSize] = { 'S', 'V', 'I', 'S', 'I', 'O', 'N', '1' };
inline bool failed(long status)
{
    return status < 0;
}

constexpr ULONG roundUpTo16(ULONG nBytes)
{
    return (nBytes + 15) / 16 * 16;
}

bool getThreadsAndBlockSize(ULONG bytes, ULONG &threads, ULONG &blockSize)
{
    constexpr ULONG kMinBlockSize = 1 << 20;
    if (bytes == 0) return false;
    ULONG nMaxThreads = std::thread::hardware_concurrency();
    threads = std::min(nMaxThreads, (bytes + kMinBlockSize - 1) / kMinBlockSize);
    blockSize = ((bytes + threads - 1) / threads + kMinBlockSize - 1) /
                kMinBlockSize * kMinBlockSize;
    threads = bytes / blockSize + !!(bytes % blockSize);
    return true;
};

// generate real AES key value from user password
// can and more operations here (such as add salt to pwd) or
// rewrite this implemention completely.
bool generateAESKeyFromUserPwd(AES::Strength strength,
                               const std::vector<byte> &pwd, std::vector<byte> &key)
{
    std::vector<byte> sha1;
    if (!AES::calculateSHA1(pwd, sha1)) return false;

    const int nKeyBytes = strength == AES::Strength::AES256 ? 32 :
                          strength == AES::Strength::AES192 ? 24 : 16;
    key.assign(nKeyBytes, 0);
    memcpy(key.data(), sha1.data(), std::min<size_t>(sha1.size(), key.size()));
    return true;
}

BCRYPT_ALG_HANDLE openAESAlg()
{
    BCRYPT_ALG_HANDLE hAES;
    auto result = BCryptOpenAlgorithmProvider(&hAES,
                                              BCRYPT_AES_ALGORITHM,
                                              nullptr,
                                              0);
    return failed(result) ? nullptr : hAES;
}

bool setChainingMode(BCRYPT_ALG_HANDLE hAES, LPCWSTR mode)
{
    auto result = BCryptSetProperty(hAES,
                                    BCRYPT_CHAINING_MODE,
                                    (PBYTE)mode,
                                    (ULONG)wcslen(mode),
                                    0);
    return !failed(result);
}

DWORD getBlockLength(BCRYPT_ALG_HANDLE hAES)
{
    DWORD blockBytes = 0, noused = 0;
    auto result = BCryptGetProperty(hAES,
                                    BCRYPT_BLOCK_LENGTH,
                                    (PBYTE)&blockBytes,
                                    sizeof(DWORD),
                                    &noused,
                                    0);
    return failed(result) ? 0 : blockBytes;
}

DWORD getKeyOBjectBytes(BCRYPT_ALG_HANDLE hAES)
{
    DWORD objectBytes = 0, noused;
    auto result = BCryptGetProperty(hAES,
                                    BCRYPT_OBJECT_LENGTH,
                                    (PBYTE)&objectBytes,
                                    sizeof(DWORD),
                                    &noused,
                                    0);
    return failed(result) ? 0 : objectBytes;
}

BCRYPT_KEY_HANDLE generateKey(BCRYPT_ALG_HANDLE hAES,
                              const std::vector<BYTE> &secret)
{
    BCRYPT_KEY_HANDLE key;
    // The memory for the key object is allocated and freed by BCrypt API, need Windows Vista with SP1 or later.
    // auto objectBytes = getKeyOBjectBytes(hAES);
    // static std::vector<BYTE> keyData(objectBytes);
    auto result = BCryptGenerateSymmetricKey(hAES,
                                             &key,
                                             nullptr, // keyData.data(),
                                             0, // objectBytes,
                                             (PUCHAR)secret.data(),
                                             (ULONG)secret.size(),
                                             0);
    return failed(result) ? nullptr : key;
}

bool destroyKey(BCRYPT_KEY_HANDLE hKey)
{
    if (!hKey) return false;
    auto result = BCryptDestroyKey(hKey);
    return !failed(result);
}

bool closeAESAlog(BCRYPT_ALG_HANDLE hAES)
{
    if (!hAES) return false;
    auto result = BCryptCloseAlgorithmProvider(hAES, 0);
    return !failed(result);
}

class RAII
{
    BCRYPT_ALG_HANDLE _hAES{};
    BCRYPT_KEY_HANDLE _hKey{};
    ULONG _blockBytes{};
public:
    RAII() = default;
    bool init(AES::Strength strength, const std::vector<byte> &pwd,
              LPCWSTR pcwChainingMode)
    {
        if (!pcwChainingMode) return false;
        _hAES = openAESAlg();
        if (!_hAES) return false;
        if (!setChainingMode(_hAES, pcwChainingMode)) return false;
        std::vector<AES::Byte> secret;
        if (!generateAESKeyFromUserPwd(strength, pwd, secret)) return false;
        _hKey = generateKey(_hAES, secret);
        if (!_hKey) return false;
        _blockBytes = getBlockLength(_hAES);
        if (!_blockBytes) return false;
        return true;
    }

    BCRYPT_KEY_HANDLE getKey() const { return _hKey; }
    ULONG getBlockBytes() const { return _blockBytes; }

    ~RAII()
    {
        destroyKey(_hKey);
        closeAESAlog(_hAES);
    }
};

bool encryptOrDecrypt(BCRYPT_KEY_HANDLE hKey,
                      DWORD blockBytes,
                      const void *buffer,
                      unsigned long bytes,
                      std::vector<BYTE> &initVector,
                      void *out,
                      ULONG outBytes,
                      bool enc)
{
    decltype(BCryptEncrypt) *op = enc ? BCryptEncrypt : BCryptDecrypt;
    ULONG targetBytes = 0;
    auto flag = (bytes % blockBytes) ? BCRYPT_BLOCK_PADDING : 0;
    BCryptEncrypt(hKey, (PBYTE)buffer, bytes, nullptr,
                  nullptr, 0, nullptr, 0, &targetBytes, flag);

    DWORD noused;
    if (outBytes < targetBytes)
        return false;
    auto result = op(hKey,
                     (PBYTE)buffer,
                     bytes,
                     nullptr,
                     initVector.empty() ?  nullptr : initVector.data(),
                     (ULONG)initVector.size(),
                     (PUCHAR)out,
                     targetBytes,
                     &noused,
                     flag);
    return !failed(result);
}

bool encryptOrDecrypt(const void *buffer,
                      unsigned long nBytes,
                      AES::Strength strength,
                      const std::vector<byte> &pwd,
                      LPCWSTR chainingMode,
                      const std::vector<byte> &initVector,
                      void *out,
                      ULONG outBytes,
                      bool enc)
{
    RAII raii;
    if (!raii.init(strength, pwd, chainingMode)) return false;
    std::vector<byte> iv(initVector);
    iv.resize(raii.getBlockBytes());
    if (wcscmp(chainingMode, BCRYPT_CHAIN_MODE_ECB) == 0)
        iv.clear();
    if (!encryptOrDecrypt(raii.getKey(), raii.getBlockBytes(), buffer, nBytes, iv,
                          out, outBytes, enc))
        return false;
    return true;
}

// Multi thread version (for ECB chaining mode only)
bool encryptOrDecryptMT(const void *buffer,
                        unsigned long bytes,
                        AES::Strength strength,
                        const std::vector<byte> &pwd,
                        void *out,
                        ULONG outBytes,
                        bool enc)
{
    if (bytes == 0) return true;
    RAII raii;
    if (!raii.init(strength, pwd, BCRYPT_CHAIN_MODE_ECB)) return false;
    ULONG targetBytes = 0;
    auto flag = (bytes % raii.getBlockBytes()) ? BCRYPT_BLOCK_PADDING : 0;
    BCryptEncrypt(raii.getKey(), (PBYTE)buffer, bytes, nullptr,
                  nullptr, 0, nullptr, 0, &targetBytes, flag);
    if (outBytes < targetBytes) return false;

    ULONG tids{}, blockSize{};
    getThreadsAndBlockSize(bytes, tids, blockSize);

    auto proc = [&raii, buffer, bytes, enc, blockSize, &out](int i) {
        ULONG nBytes = std::min(blockSize, bytes - i * blockSize);
        ULONG oBytes = roundUpTo16(nBytes);
        std::vector<AES::Byte>iv;
        if (!encryptOrDecrypt(raii.getKey(),
                              raii.getBlockBytes(),
                              (char *)buffer + i * blockSize,
                              nBytes,
                              iv,
                              (char *)out + i * blockSize,
                              oBytes,
                              enc))
            return false;
        return true;
    };

    std::vector<std::future<bool>> results(tids);
    for (ULONG i = 0; i < tids; i++)
        results[i] = std::async(std::launch::async, proc, i);

    bool ok = true;
    for (auto &f : results)
        if (!f.get()) {
            ok = false;
            break;
        }

    return ok;
}

CipherVersion version(unsigned char v)
{
    CipherVersion ver = (CipherVersion)v;
    return (ver >= CipherVersion::MiniVersion &&
            ver <= CipherVersion::CurrentVersion) ?
           ver : CipherVersion::Unknown;
}

bool isMagicHeaderValid(const void *data, CipherVersion *ver = nullptr)
{
    const auto kAllMagics = { kMagicV3, kMagicV2, kMagicV1 };
    for (const auto &e : kAllMagics)
        if (memcmp(data, e, kMagicSize) == 0) {
            if (ver) *ver = version(static_cast<const char *>(data)[kMagicSize - 1] - '0');
            return true;
        }
    return false;
}

CipherVersion cipheredFileVersion(const QString &file)
{
    const auto &localFile = file.toLocal8Bit();
    std::ifstream ifs(localFile, std::ifstream::binary | std::ifstream::ate);
    if (!ifs.is_open()) return CipherVersion::Unknown;
    unsigned long fileSize = (unsigned long)ifs.tellg();
    if (fileSize <= kHeaderSize) return CipherVersion::Unknown;
    char data[kMagicSize] = {};
    ifs.seekg(0, std::ifstream::beg);
    ifs.read(data, kMagicSize);
    ifs.close();
    if (ifs.bad()) return CipherVersion::Unknown;
    return version(data[kMagicSize - 1] - '0');
}

CipherVersion cipheredBlockVersion(const void *block)
{
    if (!block || !isMagicHeaderValid(block)) return CipherVersion::Unknown;
    return version(((char *)block)[kMagicSize - 1] - '0');
}

std::vector<AES::Byte> defaultKey(CipherVersion ver)
{
    auto pwd = ver == CipherVersion::Version1 ? AES::uniformCodeV1() :
               ver == CipherVersion::Version2 ? AES::uniformCodeV2() :
               ver == CipherVersion::CurrentVersion ? AES::uniformCodeV3() :
               std::vector<AES::Byte>();
    return pwd;
}

std::vector<AES::Byte> defaultKeyFromFile(const QString &file)
{
    return defaultKey(cipheredFileVersion(file));
}

std::vector<AES::Byte> defaultKeyFromBlock(const void *block)
{
    return defaultKey(cipheredBlockVersion(block));
}

} // namespace

namespace AES
{
std::vector<AES::Byte> uniformCodeForAuthority()
{
    return std::vector<Byte> { 0x02, 0x00, 0x02, 0x00, 0x00, 0x05, 0x00, 0x07 };
}

std::vector<AES::Byte> uniformCodeV3()
{
    return std::vector<Byte> { 0x02, 0x00, 0x02, 0x01, 0x01, 0x02, 0x02, 0x01 };
}

std::vector<AES::Byte> uniformCodeV2()
{
    return std::vector<Byte> { 0x02, 0x00, 0x01, 0x09, 0x00, 0x02, 0x02, 0x05 };
}

std::vector<Byte> uniformCodeV1()
{
    return std::vector<Byte> { 0x02, 0x00, 0x01, 0x08, 0x00, 0x03, 0x02, 0x07 };
}

std::vector<AES::Byte> keyGeneratorCode()
{
    const char code[] = "KeyGeneratorCode";
    return std::vector<Byte>(code, code + sizeof(code));
}

bool calculateSHA1(const std::vector<Byte> &text, std::vector<Byte> &result)
{
    result.clear();
    BCRYPT_ALG_HANDLE hSHA1;
    auto algoID = BCRYPT_SHA1_ALGORITHM;
    if (failed(BCryptOpenAlgorithmProvider(&hSHA1, algoID, nullptr, 0)))
        return false;
    DWORD cbHashObject, noused;
    if (failed(BCryptGetProperty(hSHA1, BCRYPT_OBJECT_LENGTH, (PBYTE)&cbHashObject,
                                 sizeof(DWORD), &noused, 0)))
        return false;
    BCRYPT_HASH_HANDLE hash;
    std::vector<BYTE> hashObj(cbHashObject);
    if (failed(BCryptCreateHash(hSHA1, &hash, hashObj.data(), cbHashObject,
                                nullptr, 0, 0)))
        return false;
    if (failed(BCryptHashData(hash, const_cast<PUCHAR>(text.data()),
                              (ULONG)text.size(), 0)))
        return false;

    result.resize(20);
    if (failed(BCryptFinishHash(hash, result.data(), (ULONG)result.size(), 0)))
        return false;
    BCryptDestroyHash(hash);
    BCryptCloseAlgorithmProvider(hSHA1, 0);

    return true;
}

bool encryptFile(const QString &file, const QString &targetFile,
                 const std::vector<Byte> &pwd)
{
    const auto &localFile = file.toLocal8Bit();
    std::ifstream ifs(localFile, std::ifstream::binary | std::ifstream::ate);
    if (!ifs.is_open()) return false;
    unsigned long fileBytes = (unsigned long)ifs.tellg();
    std::vector<Byte> fileData(fileBytes);
    ifs.seekg(0, std::ifstream::beg);
    ifs.read((char *)fileData.data(), fileData.size());
    ifs.close();

    CipheredFileWriter fw(targetFile, pwd);
    fw.writeData(fileData);
    fw.writeVerification();
    fw.close();
    return true;
}

bool decryptFile(const QString &file, const QString &targetFile,
                 const std::vector<Byte> &pwd)
{
    auto innerPwd = pwd.empty() ? defaultKeyFromFile(file) : pwd;
    CipheredFileReader fr(file, innerPwd);
    if (!fr.isValid()) return false;
    std::vector<Byte> fileData;
    if (!fr.readAll(fileData)) return false;
    fr.close();

    const auto &localTargetFile = targetFile.toLocal8Bit();
    std::ofstream ofs(localTargetFile, std::ifstream::binary);
    if (!ofs.is_open()) return false;
    ofs.write((char *)fileData.data(), fileData.size());
    if (ofs.bad()) return false;
    ofs.close();
    return true;
}

bool mayBeCipher(const QString &file)
{
    const auto &localFile = file.toLocal8Bit();
    std::ifstream ifs(localFile, std::ifstream::binary | std::ifstream::ate);
    if (!ifs.is_open()) return false;
    unsigned long fileSize = (unsigned long)ifs.tellg();
    if (fileSize <= kHeaderSize) return false;
    char fileData[kMagicSize] = {};
    ifs.seekg(0, std::ifstream::beg);
    ifs.read(fileData, kMagicSize);
    ifs.close();
    return isMagicHeaderValid(fileData);
}

bool mayBeCipher(const void *buf, int nBytes)
{
    if (nBytes <= kHeaderSize) return false;
    return isMagicHeaderValid(buf);
}

void Encryptor::setInitialVector(const std::vector<Byte> &iv)
{
    _iv = iv;
}

std::vector<Byte> Encryptor::initialVector() const
{
    return _iv;
}

void Encryptor::setPlaintTextBytes(unsigned long bytes)
{
    _plainTextBytes = bytes;
}

void Encryptor::setCBCChainingMode(bool on)
{
    _useCBCChainingMode = on;
}

Encryptor::Encryptor(Strength s, const std::vector<Byte> &iv)
    : _strength(s), _iv(iv)
{
}

void Encryptor::setStrength(Strength s)
{
    _strength = s;
}

Strength Encryptor::strength() const
{
    return _strength;
}

bool Encryptor::encrypt(const void *buffer, unsigned long nBytes,
                        const std::vector<Byte> &pwd, std::vector<Byte> &out)
{
    out.resize(roundUpTo16(nBytes));
    return _useCBCChainingMode ?
           encryptOrDecrypt(buffer, nBytes, _strength, pwd, BCRYPT_CHAIN_MODE_CBC, _iv,
                            out.data(), (ULONG)out.size(), true) :
           encryptOrDecryptMT(buffer, nBytes, _strength, pwd, out.data(),
                              (ULONG)out.size(), true);
}

bool Encryptor::encrypt(const void *buffer, unsigned long nBytes,
                        const std::vector<Byte> &pwd, void *out, unsigned long outBytes)
{
    return _useCBCChainingMode ?
           encryptOrDecrypt(buffer, nBytes, _strength, pwd, BCRYPT_CHAIN_MODE_CBC, _iv,
                            out, outBytes, true) :
           encryptOrDecryptMT(buffer, nBytes, _strength, pwd, out, outBytes, true);
}

bool Encryptor::decrypt(const void *buffer, unsigned long nBytes,
                        const std::vector<Byte> &pwd, std::vector<Byte> &out, bool resetLength)
{
    out.resize(roundUpTo16(nBytes));
    bool ok = _useCBCChainingMode ?
              encryptOrDecrypt(buffer, nBytes, _strength, pwd, BCRYPT_CHAIN_MODE_CBC, _iv,
                               out.data(), (ULONG)out.size(), false) :
              encryptOrDecryptMT(buffer, nBytes, _strength, pwd, out.data(),
                                 (ULONG)out.size(), false);

    if (!ok) return false;
    if (resetLength && _plainTextBytes) out.resize(_plainTextBytes);
    return true;
}

bool Encryptor::decrypt(const void *buffer, unsigned long nBytes,
                        const std::vector<Byte> &pwd,
                        void *out, unsigned long outBytes, unsigned long *originBytes)
{
    bool ok = _useCBCChainingMode ?
              encryptOrDecrypt(buffer, nBytes, _strength, pwd, BCRYPT_CHAIN_MODE_CBC, _iv,
                               out, outBytes, false) :
              encryptOrDecryptMT(buffer, nBytes, _strength, pwd, out, outBytes, false);

    if (!ok) return false;
    if (_plainTextBytes && originBytes) *originBytes = _plainTextBytes;
    return true;
}

struct CipheredFileWriter::Impl {
    std::fstream fs;
    std::vector<Byte> pwd;
};

CipheredFileWriter::CipheredFileWriter(const QString &file,
                                       const std::vector<Byte> &pwd)
    : _pImpl(new Impl)
{
    _pImpl->pwd = pwd;
    const auto &localFile = file.toLocal8Bit();
    std::ofstream of(localFile, std::fstream::out | std::fstream::trunc);
    of.close();
    _pImpl->fs.open(localFile, std::fstream::in | std::fstream::out |
                    std::fstream::binary);
}

CipheredFileWriter::~CipheredFileWriter()
{
    close();
}

bool CipheredFileWriter::writeData(const void *buffer, unsigned long bytes)
{
    Encryptor enc(Strength::AES256);
    enc.setCBCChainingMode(false);
    std::vector<Byte> out;
    if (!enc.encrypt(buffer, bytes, _pImpl->pwd, out)) return false;
    if (!_pImpl->fs.is_open()) return false;
    _pImpl->fs.seekp(0, std::ios_base::end);
    if (!_pImpl->fs.tellp()) {
        char header[kHeaderSize] = {};
        memcpy(header, kMagicV3, kMagicSize);
        if (!_pImpl->fs.write(header, sizeof(header))) return false;
    }

    if (!_pImpl->fs.write((char *)&bytes, sizeof(bytes))) return false;
    unsigned int blockBytes = (unsigned int)out.size();
    if (!_pImpl->fs.write((char *)&blockBytes, sizeof(blockBytes))) return false;
    if (!_pImpl->fs.write((char *)out.data(), out.size())) return false;
    _pImpl->fs.flush();
    return true;
}

bool CipheredFileWriter::writeData(const std::vector<Byte> &data)
{
    return writeData((void *)data.data(), (unsigned long)data.size());
}

bool CipheredFileWriter::writeVerification()
{
    if (!_pImpl->fs.is_open()) return false;
    auto cp = _pImpl->fs.tellp();
    _pImpl->fs.seekp(0, std::ios_base::end);
    auto ep = _pImpl->fs.tellp();
    if (ep <= kHeaderSize) return false;
    _pImpl->fs.seekg(0, std::ios_base::beg);
    char header[kHeaderSize] = {};
    if (!_pImpl->fs.read(header, sizeof(header))) return false;

    std::vector<Byte> data((unsigned long)ep - kHeaderSize);
    if (!_pImpl->fs.read((char *)data.data(), data.size())) return false;
    std::vector<Byte> sha1;
    if (!AES::calculateSHA1(data, sha1)) return false;

    int verificationFlag = 1;
    memcpy(header + kMagicSize,  &verificationFlag,
           sizeof(verificationFlag));
    memcpy(header + kMagicSize + 4, sha1.data(), sha1.size());
    _pImpl->fs.seekp(0, std::ios_base::beg);
    if (!_pImpl->fs.write(header, sizeof(header))) return false;
    _pImpl->fs.flush();
    _pImpl->fs.seekp(cp, std::ios_base::beg);
    return true;
}

void CipheredFileWriter::close()
{
    if (_pImpl->fs.is_open())
        _pImpl->fs.close();
}

struct CipheredFileReader::Impl {
    std::ifstream ifs;
    std::vector<Byte> pwd;
    CipherVersion ver = CipherVersion::Unknown;
};

CipheredFileReader::CipheredFileReader(const QString &file,
                                       const std::vector<Byte> &pwd)
    : _pImpl(new Impl)
{
    _pImpl->pwd = pwd.empty() ? defaultKeyFromFile(file) : pwd;
    const auto &localFile = file.toLocal8Bit();
    _pImpl->ifs.open(localFile, std::fstream::in | std::fstream::binary);
}

CipheredFileReader::~CipheredFileReader()
{
    close();
}

bool CipheredFileReader::isValid()
{
    if (!_pImpl->ifs.is_open()) return false;
    _pImpl->ifs.seekg(0, std::fstream::end);
    auto tg = _pImpl->ifs.tellg();
    if (tg <= kHeaderSize) return false;
    _pImpl->ifs.seekg(0, std::fstream::beg);
    char header[kHeaderSize] = {};
    if (!_pImpl->ifs.read(header, sizeof(header))) return false;
    if (!isMagicHeaderValid(header, &_pImpl->ver)) return false;
    int verificatioinFlag = *(int *)((char *)header + kMagicSize);
    if (verificatioinFlag != 0 && verificatioinFlag != 1) return false;
    if (verificatioinFlag == 1) {
        std::vector<Byte> raw((unsigned long)tg - kHeaderSize);
        if (!_pImpl->ifs.read((char *)raw.data(), raw.size())) return false;
        std::vector<Byte> sha1;
        AES::calculateSHA1(raw, sha1);
        if (memcmp(sha1.data(), header + kMagicSize + 4,
                   sha1.size()))  return false;
    }
    return true;
}

bool CipheredFileReader::readAll(std::vector<Byte> &data)
{
    if (!isValid()) return false;
    data.clear();
    _pImpl->ifs.seekg(kHeaderSize, std::ifstream::beg);
    while (true) {
        unsigned long plainBytes = 0u;
        if (!_pImpl->ifs.read((char *)&plainBytes, sizeof(plainBytes))) {
            if (_pImpl->ifs.eof()) break;
            return false;
        }
        unsigned long blockBytes = 0u;
        if (!_pImpl->ifs.read((char *)&blockBytes, sizeof(blockBytes))) return false;

        std::vector<Byte> block(blockBytes);
        if (!_pImpl->ifs.read((char *)block.data(), block.size())) return false;

        Encryptor enc(Strength::AES256);
        enc.setPlaintTextBytes(plainBytes);
        enc.setCBCChainingMode(_pImpl->ver < CipherVersion::Version3);

        std::vector<Byte> blockPlain;
        if (!enc.decrypt(block.data(), (unsigned long)block.size(),
                         _pImpl->pwd, blockPlain)) return false;
        data.resize(data.size() + blockPlain.size());
        memcpy(data.data() + data.size() - blockPlain.size(), blockPlain.data(),
               blockPlain.size());
    }

    return true;
}

void CipheredFileReader::close()
{
    if (_pImpl->ifs.is_open())
        _pImpl->ifs.close();
}

int CipheredMemBlock::fixedSize()
{
    return kHeaderSize + sizeof(int) + sizeof(int);
}

CipheredMemBlock::CipheredMemBlock(const void *addr, uint32_t bytes)
    : _memBlock(bytes)
{
    memcpy(_memBlock.data(), addr, bytes);
}

std::unique_ptr<CipheredMemBlock> CipheredMemBlock::fromData(const void *data,
        uint32_t dataBytes,
        const std::vector<Byte> &pwd,
        bool chainingMode)
{
    if (!data) return nullptr;
    Encryptor enc(Strength::AES256);
    enc.setCBCChainingMode(chainingMode);
    auto innerPwd = pwd.empty() ? AES::uniformCodeV3() : pwd;
    std::unique_ptr<CipheredMemBlock> ptrBlock(
        std::make_unique<CipheredMemBlock>());
    uint outSize = roundUpTo16(dataBytes);
    ptrBlock->_memBlock.resize(fixedSize() + outSize);
    if (!enc.encrypt(data, dataBytes, innerPwd,
                     ptrBlock->_memBlock.data() + fixedSize(), outSize))
        return nullptr;

    memcpy(ptrBlock->_memBlock.data(),  chainingMode ? kMagicV2 : kMagicV3,
           kMagicSize);
    memcpy(ptrBlock->_memBlock.data() + kHeaderSize, &dataBytes, sizeof(int));
    memcpy(ptrBlock->_memBlock.data() + kHeaderSize + sizeof(int), &outSize,
           sizeof(int));
    return ptrBlock;
}

bool CipheredMemBlock::toData(std::vector<Byte> &data,
                              const std::vector<Byte> &pwd)
{
    data.clear();
    if (!isValid()) return false;
    uint plainBytes = getPlainTextBytes();
    if (plainBytes == 0u)
        return true;
    data.resize(roundUpTo16(plainBytes));
    uint bytesCopied{};
    if (!toData(data.data(), (uint)data.size(), pwd, &bytesCopied)) return false;
    data.resize(bytesCopied);
    return true;
}

bool CipheredMemBlock::toData(void *data, uint bytes,
                              const std::vector<Byte> &pwd, uint *pBytesCopied)
{
    if (!data) return false;
    if (!isValid()) return false;
    int ver = getVersion();
    uint plainBytes = getPlainTextBytes();
    uint cipherBytes = getCipherTextBytes();
    Encryptor enc(Strength::AES256);
    enc.setPlaintTextBytes(plainBytes);
    enc.setCBCChainingMode(ver < (int)CipherVersion::Version3);
    auto innerPwd = pwd.empty() ? defaultKeyFromBlock(_memBlock.data()) : pwd;
    if (!enc.decrypt(_memBlock.data() + fixedSize(), (unsigned long)cipherBytes,
                     innerPwd, data, bytes)) return false;
    if (pBytesCopied) *pBytesCopied = plainBytes;
    return true;
}

bool CipheredMemBlock::toDataFirstNBytes(uint bytes, std::vector<Byte> &data,
                                         const std::vector<Byte> &pwd)
{
    data.clear();
    if (!isValid()) return false;
    int ver = getVersion();
    uint cipherBytes = getCipherTextBytes();
    bytes = std::min(bytes, cipherBytes);
    Encryptor enc(Strength::AES256);
    enc.setCBCChainingMode(ver < (int)CipherVersion::Version3);
    auto innerPwd = pwd.empty() ? defaultKeyFromBlock(_memBlock.data()) : pwd;
    if (!enc.decrypt(_memBlock.data() + fixedSize(), bytes, innerPwd, data,
                     false)) return false;
    return true;
}

std::vector<Byte> &CipheredMemBlock::memBlock()
{
    return _memBlock;
}

bool CipheredMemBlock::isValid() const
{
    if (_memBlock.size() < fixedSize()) return false;
    if (!isMagicHeaderValid(_memBlock.data())) return false;
    uint cipherBytes = *(uint *)(_memBlock.data() + kHeaderSize + sizeof(uint));
    if (fixedSize() + cipherBytes != _memBlock.size()) return false;
    return true;
}

int CipheredMemBlock::getVersion() const
{
    if (!isValid()) return (int)CipherVersion::Unknown;
    int ver = (int)version(_memBlock[kMagicSize - 1] - '0');
    return ver;
}

uint CipheredMemBlock::getPlainTextBytes() const
{
    if (!isValid()) return 0;
    uint bytes = *(uint *)(_memBlock.data() + kHeaderSize);
    return bytes;
}

uint CipheredMemBlock::getCipherTextBytes() const
{
    if (!isValid()) return 0;
    uint cipherBytes = *(uint *)(_memBlock.data() + kHeaderSize + sizeof(uint));
    return cipherBytes;
}
} // namespace AES
